"""
# -*- coding: utf-8 -*-
# @Time    : 2023/5/24 20:36
# @Author  : 王摇摆
# @FileName: AVL.py
# @Software: PyCharm
# @Blog    ：https://blog.csdn.net/weixin_44943389?type=blog
"""
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

from ANN.Model_Manual import distance


def means_step(X, iteration_steps=200):
    count = X.shape[0]
    i = np.random.randint(0, count)
    j = np.random.randint(0, count - 1)
    j += (j >= i)
    ic = 1
    jc = 1
    p = X[i]
    q = X[j]
    ps = np.zeros((iteration_steps, len(p)))
    qs = np.zeros((iteration_steps, len(q)))
    ts = np.zeros((iteration_steps, len(q)))
    for l in range(iteration_steps):
        p, q, ic, jc, _ = step(X, p, q, ic, jc)
        ps[l], qs[l], ts[l] = p, q, _
    return ps, qs, ts


def step(X, p, q, ic, jc):
    k = np.random.randint(0, X.shape[0])
    di = ic * distance(p, X[k])
    dj = jc * distance(q, X[k])
    if di == dj:
        return
    if di < dj:
        p = (p * ic + X[k]) / (ic + 1)
        ic = ic + 1
    else:
        q = (q * jc + X[k]) / (jc + 1)
        jc = jc + 1
    return p, q, ic, jc, X[k]


fig, ax = plt.subplots()
ax.set_facecolor('#f8f9fa')

np.random.seed(0)
X = np.array([[1, 1], [2, 3], [3, -1], [0, 0], [-1, -2], [-2, 2], [-4, 4], [3, -1], [1, -4], [0, 2], [-3, 0]])
itr = 50
ps, qs, ts = means_step(X, iteration_steps=itr)
x_min, x_max = X[:, 0].min() - .5, X[:, 0].max() + .5
y_min, y_max = X[:, 1].min() - .5, X[:, 1].max() + .5
plt.xlim(x_min, x_max)
plt.ylim(y_min, y_max)
x1 = X[:, 0]
y1 = X[:, 1]
plt.scatter(x1, y1, c='#e63946', marker='.')

p = plt.scatter(ps[0][0], ps[0][1], c='#457b9d', marker='$P$', s=60)
q = plt.scatter(qs[0][0], qs[0][1], c='#457b9d', marker='$Q$', s=60)
t = plt.scatter(ts[0][0], ts[0][1], c='#457b9d', marker='$T$', s=60)


def update(i):
    p.set_offsets(ps[i])
    q.set_offsets(qs[i])
    if i + 1 < itr:
        t.set_offsets(ts[i + 1])
    else:
        t.set_offsets([10, 10])
    return p, q, t


ani = animation.FuncAnimation(fig, update, range(1, itr), interval=500, blit=True, repeat=False)
ani.save('means.gif')
plt.show()
